from PyEMD import EEMD
import numpy as np
import matplotlib.pyplot as plt

# EEDM模态分解
def eedm(data, n):
    '''
    n: 分解次数
    '''
    S = np.array(data)
    # EEMD计算
    eemd = EEMD()
    eemd.trials = 50
    eemd.noise_seed(12345)
    E_IMFs = eemd.eemd(S, max_imf=n)
    imfNo = E_IMFs.shape[0]
    return E_IMFs

def plot_imfs(imfs):
    # plot results
    for i in range(imfs.shape[0]):
        _ = plt.plot(imfs[i], 'g')
        _ = plt.title('Power of ' + 'Imf'+str(i+1)+' component')
        # _ = plt.savefig(work_dir + '/../Image/EEMD模态分解.svg', format='svg')
        _ = plt.show()